
import JT
import time
import numpy as np

# Define which data to use for registration:
vert = 1
proj = 1
body = 4

Data = JT.Data.gold_standard()    
#===============================================================================
#    Setup DRR
#===============================================================================      
drr = JT.DRR()
drr.SetExternalCalibration(Data.Frames[vert].FixedImages[proj].calibration)

drr.AddVolume(Data.Frames[vert].Volumes[0].data)
drr.AddBodyCenter(Data.Frames[vert].Volumes[0].BodyCenters[body])
drr._renWin.SetOffScreenRendering(0)

reg=JT.Registration()
reg.SetFixedImage(Data.Frames[vert].FixedImages[proj].data.GetImage(0))
reg.SetMovingImageGenerator(drr)
reg.SetRegionOfInterest(Data.Frames[vert].FixedImages[proj].ROIs[body].data)

#===============================================================================
#    Setup Registration
#===============================================================================      
# set Metric
reg.SetImageMetric('bpc')
# Set Optimizer
gd = JT.RegularStepGradientDescentOptimizer()
reg.SetOptimizer(gd)

Testing = JT.Testing()
 # Set the starting pose
pose = Testing.GetStartPose(150)
print "Starting Pose: ", pose
reg.SetInitialPose(pose)

# Set Scales 1/1000 for translations, 1 for rotations
scales = [1.0,1.0,1.0,1700.0,1700.0,1700.0]
reg.SetScales(scales)

# Registration Options
reg._plot = True
reg._debug = True
reg._logGradientResult = False
gd._debug = False

#===============================================================================
# Setup Optimizer Parameters
#===============================================================================
# Set the cost function
gd.SetCostFunction(reg)

# Modify Optimizer Parameters
gd.SetMaximumStepLength(4)
gd.SetMinimumStepLength(2)
gd.SetRelaxationFactor(0.85)
gd.SetNumberOfIterations(50)
gd._MinimumGradientStepLength = 0.25

def tuneMaximumStepLength():    
    # For pose 50: looks like 3.5-4.0 is best for translations, and 4.5 or
    #  greater is best for rotations
    result = {}
    parameter_range = [2.0,3.0,4.0,5.0,6.0,8.0]
    for par in parameter_range:
        gd.SetMaximumStepLength(par)
        gd._MaximumGradientStepLength=par
        reg.Start()
        # Print Final State    
        gd.PrintSelf() 
        result[par] = reg._BestPose 
        reg.Reset()
    
    print result   

def tuneRelaxationFactor():    
    # looks like 0.75 is a good choice
    result = {}
    parameter_range = [0.25,0.35,0.45,0.55,0.65,0.75,0.85,0.95]
    for par in parameter_range:
        gd.SetRelaxationFactor(par)
        reg.Start()
        # Print Final State    
        gd.PrintSelf() 
        result[par] = reg._BestPose 
        reg.Reset()
    
    print result     
    
def tuneSensitivity():
    pose = Testing.GetStartPose(0)
    parameter_range = [1e-3,1e-2,1e-1,5e-1,1,2,5]
    result = {}
    for par in parameter_range:
        result[par] = reg.GetGradient(pose,par)
        g = np.asarray(result[par])
        trans = np.sqrt((g[0:3]**2).sum())
        rot = np.sqrt((g[3:6]**2).sum())
        print par,": t=",trans,", r=",rot,", ratio=",rot/trans
    print result

def tuneScales():    
    # looks like 0.75 is a good choice
    result = {}
    parameter_range = [1900]
    for par in parameter_range:
        gd._MaximumGradientStepLength = 6
        gd.SetMaximumStepLength(15)
        scales = [1,1,1,par,par,par]
        gd.SetScales(scales)
        reg.Start()
        # Print Final State    
        gd.PrintSelf() 
        result[par] = reg.bestPose 
        reg.Reset()
    print result
  

if __name__ == "__main__":
    #T = Test(logging=False,upload=False)
    #tuneRelaxationFactor()
    #tuneMaximumStepLength()
    #tuneSensitivity()
    tuneScales()
    #time.sleep(7)
    #JT.PlotResultsFromFile.main(0)
    
